// csp.go
package main

import (
	"fmt"
	"math/big"

	"github.com/Nik-U/pbc"
)

// CSP 代表云服务商
type CSP struct {
	Name    string
	sk      *pbc.Element      // 私钥 αx
	Pk      *pbc.Element      // 公钥 pkx = h^αx
	storage map[string][]byte // 内存存储，替代文件系统
}

// Proof 包含了CSP为一次审计生成的所有证明组件
type Proof struct {
	Theta *pbc.Element
	Beta  *pbc.Element
	Psi   map[string]*pbc.Element
}

// NewCSP 创建一个新的CSP实例
func NewCSP(name string) *CSP {
	return &CSP{
		Name:    name,
		storage: make(map[string][]byte),
	}
}

// GenerateKeyPair 为CSP生成公私钥对
func (c *CSP) GenerateKeyPair(sys *System) {
	c.sk = sys.PublicPar.Pairing.NewZr().Rand()
	c.Pk = sys.PublicPar.Pairing.NewG1().PowZn(sys.PublicPar.H, c.sk)
}

// StoreData 将数据存储在内存中
func (c *CSP) StoreData(dataID string, data []byte) {
	// 创建一个副本以模拟真实存储，避免多个CSP共享同一个内存地址
	dataCopy := make([]byte, len(data))
	copy(dataCopy, data)
	c.storage[dataID] = dataCopy
}

// splitDataIntoBlocks 内存中的数据分块
func (c *CSP) splitDataIntoBlocks(data []byte, blockSize int, pairing *pbc.Pairing) []*pbc.Element {
	var blocks []*pbc.Element
	numBlocks := len(data) / blockSize
	if len(data)%blockSize != 0 {
		numBlocks++
	}

	for i := 0; i < numBlocks; i++ {
		start := i * blockSize
		end := start + blockSize
		if end > len(data) {
			end = len(data)
		}

		// 1. 将字节切片转换为大整数
		mi_bigint := new(big.Int).SetBytes(data[start:end])

		// 2. 将大整数设置到Zr元素中，这会自动进行取模操作
		blockElement := pairing.NewZr().SetBig(mi_bigint)

		blocks = append(blocks, blockElement)
	}
	return blocks
}

// GenerateProof 为指定的数据和挑战生成一致性证明
func (c *CSP) GenerateProof(dataID string, tau string, sys *System, allCSPs []*CSP) (*Proof, error) {
	data, ok := c.storage[dataID]
	if !ok {
		return nil, fmt.Errorf("data with ID '%s' not found in CSP %s", dataID, c.Name)
	}

	blocks := c.splitDataIntoBlocks(data, sys.BlockSize, sys.PublicPar.Pairing)
	n := len(blocks)
	if n == 0 {
		return nil, fmt.Errorf("no data blocks found for ID '%s'", dataID)
	}

	// 1. 生成挑战 v_i = F(τ||i)
	v := make([]*pbc.Element, n)
	for i := 0; i < n; i++ {
		input := []byte(fmt.Sprintf("%s||%d", tau, i))
		v[i] = F(input, sys.PublicPar.Pairing)
	}

	// 2. 计算聚合值 sum(m_i * v_i)
	sumMV := sys.PublicPar.Pairing.NewZr().Set0()
	for i := 0; i < n; i++ {
		term := sys.PublicPar.Pairing.NewZr().Mul(blocks[i], v[i])
		sumMV.Add(sumMV, term)
	}

	// 3. 引入 d_x1, d_x2 来增强隐私性
	d1 := sys.PublicPar.Pairing.NewZr().Rand()
	d2 := sys.PublicPar.Pairing.NewZr().Rand()

	// 4. 计算证明组件 θx 和 βx (修正后的核心逻辑)
	// θx = d_{x,1} (现在只包含一个随机数)
	theta := d1

	// βx = g^(α_x * Σ(m_i*v_i) + d_{x,2})
	// 这个构造直接将私钥和数据聚合值绑定在一起
	betaExp := sys.PublicPar.Pairing.NewZr().Mul(c.sk, sumMV)
	betaExp.Add(betaExp, d2)
	beta := sys.PublicPar.Pairing.NewG1().PowZn(sys.PublicPar.G, betaExp)

	// 5. 计算 ψx,y = pk_y^dx
	d := sys.PublicPar.Pairing.NewZr().Add(d1, d2) // d = d1 + d2
	psi := make(map[string]*pbc.Element)
	for _, otherCSP := range allCSPs {
		if otherCSP.Name != c.Name {
			psi[otherCSP.Name] = sys.PublicPar.Pairing.NewG1().PowZn(otherCSP.Pk, d)
		}
	}

	return &Proof{Theta: theta, Beta: beta, Psi: psi}, nil
}
